/*
 * Copyright 2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.alibaba.cloud.ai.manus.tool.searchAPI;

import com.alibaba.cloud.ai.manus.tool.AbstractBaseTool;
import com.alibaba.cloud.ai.manus.tool.code.ToolExecuteResult;
import com.alibaba.cloud.ai.manus.tool.searchAPI.serpapi.SerpApiProperties;
import com.alibaba.cloud.ai.manus.tool.searchAPI.serpapi.SerpApiService;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.core.type.TypeReference;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.springframework.ai.openai.api.OpenAiApi;

public class GoogleSearch extends AbstractBaseTool<GoogleSearch.GoogleSearchInput> {

  private static final Logger log = LoggerFactory.getLogger(GoogleSearch.class);

  private SerpApiService service;

  private final ObjectMapper objectMapper;

  private static String PARAMETERS = """
    {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "(required) The search query to submit to Google."
            },
            "num_results": {
                "type": "integer",
                "description": "(optional) The number of search results to return. Default is 10.",
                "default": 10
            }
        },
        "required": ["query"]
    }
    """;

  private static final String name = "google_search";

  private static final String description = """
    Perform a Google search and return a list of relevant links.
    Use this tool when you need to find information on the web, get up-to-date data, or research specific topics.
    The tool returns a list of URLs that match the search query.
    """;

  public static OpenAiApi.FunctionTool getToolDefinition() {
    OpenAiApi.FunctionTool.Function function = new OpenAiApi.FunctionTool.Function(description, name, PARAMETERS);
    OpenAiApi.FunctionTool functionTool = new OpenAiApi.FunctionTool(function);
    return functionTool;
  }

  private static final String SERP_API_KEY = System.getenv("SERP_API_KEY");

  private String lastQuery = "";

  private String lastSearchResults = "";

  private Integer lastNumResults = 0;

  public GoogleSearch(ObjectMapper objectMapper) {
    this.objectMapper = objectMapper;
    service = new SerpApiService(new SerpApiProperties(SERP_API_KEY, "google"));
  }

  @SuppressWarnings("unchecked")
  public ToolExecuteResult run(String toolInput) {
    log.info("GoogleSearch toolInput:{}", toolInput);

    // Add exception handling for JSON deserialization
    try {
      Map<String, Object> toolInputMap = objectMapper.readValue(toolInput,
        new TypeReference<Map<String, Object>>() {
        });
      String query = (String) toolInputMap.get("query");
      this.lastQuery = query;

      Integer numResults = 2;
      if (toolInputMap.get("num_results") != null) {
        numResults = (Integer) toolInputMap.get("num_results");
      }
      this.lastNumResults = numResults;

      SerpApiService.Request request = new SerpApiService.Request(query);
      Map<String, Object> response = service.apply(request);

      if (response.containsKey("answer_box") && response.get("answer_box") instanceof List) {
        response.put("answer_box", ((List<Map<String, Object>>) response.get("answer_box")).get(0));
      }

      String toret = "";
      if (response.containsKey("answer_box")
        && ((Map<String, Object>) response.get("answer_box")).containsKey("answer")) {
        toret = ((Map<String, Object>) response.get("answer_box")).get("answer").toString();
      } else if (response.containsKey("answer_box")
        && ((Map<String, Object>) response.get("answer_box")).containsKey("snippet")) {
        toret = ((Map<String, Object>) response.get("answer_box")).get("snippet").toString();
      } else if (response.containsKey("answer_box")
        && ((Map<String, Object>) response.get("answer_box")).containsKey("snippet_highlighted_words")) {
        toret = ((List<String>) ((Map<String, Object>) response.get("answer_box"))
          .get("snippet_highlighted_words")).get(0);
      } else if (response.containsKey("sports_results")
        && ((Map<String, Object>) response.get("sports_results")).containsKey("game_spotlight")) {
        toret = ((Map<String, Object>) response.get("sports_results")).get("game_spotlight").toString();
      } else if (response.containsKey("shopping_results")
        && ((List<Map<String, Object>>) response.get("shopping_results")).get(0).containsKey("title")) {
        List<Map<String, Object>> shoppingResults = (List<Map<String, Object>>) response
          .get("shopping_results");
        List<Map<String, Object>> subList = shoppingResults.subList(0, 3);
        toret = subList.toString();
      } else if (response.containsKey("knowledge_graph")
        && ((Map<String, Object>) response.get("knowledge_graph")).containsKey("description")) {
        toret = ((Map<String, Object>) response.get("knowledge_graph")).get("description").toString();
      } else if ((((List<Map<String, Object>>) response.get("organic_results")).get(0)).containsKey("snippet")) {
        toret = (((List<Map<String, Object>>) response.get("organic_results")).get(0)).get("snippet")
          .toString();
      } else if ((((List<Map<String, Object>>) response.get("organic_results")).get(0)).containsKey("link")) {
        toret = (((List<Map<String, Object>>) response.get("organic_results")).get(0)).get("link").toString();
      } else if (response.containsKey("images_results")
        && ((Map<String, Object>) ((List<Map<String, Object>>) response.get("images_results")).get(0))
        .containsKey("thumbnail")) {
        List<String> thumbnails = new ArrayList<>();
        List<Map<String, Object>> imageResults = (List<Map<String, Object>>) response.get("images_results");
        for (Map<String, Object> item : imageResults.subList(0, 10)) {
          thumbnails.add(item.get("thumbnail").toString());
        }
        toret = thumbnails.toString();
      } else {
        toret = "No good search result found";
      }
      log.warn("SerpapiTool result:{}", toret);
      this.lastSearchResults = toret;
      return new ToolExecuteResult(toret);
    } catch (Exception e) {
      log.error("Error deserializing JSON", e);
      return new ToolExecuteResult("Error deserializing JSON: " + e.getMessage());
    }
  }

  @Override
  public String getName() {
    return name;
  }

  @Override
  public String getDescription() {
    return description;
  }

  @Override
  public String getParameters() {
    return PARAMETERS;
  }

  @Override
  public Class<GoogleSearchInput> getInputType() {
    return GoogleSearchInput.class;
  }

  @SuppressWarnings("unchecked")
  @Override
  public ToolExecuteResult run(GoogleSearchInput input) {
    String query = input.getQuery();
    Integer numResults = input.getNumResults() != null ? input.getNumResults() : 2;

    log.info("GoogleSearch input: query={}, numResults={}", query, numResults);

    this.lastQuery = query;
    this.lastNumResults = numResults;

    try {
      SerpApiService.Request request = new SerpApiService.Request(query);
      Map<String, Object> response = service.apply(request);

      if (response.containsKey("answer_box") && response.get("answer_box") instanceof List) {
        response.put("answer_box", ((List<Map<String, Object>>) response.get("answer_box")).get(0));
      }

      String toret = "";
      if (response.containsKey("answer_box")
        && ((Map<String, Object>) response.get("answer_box")).containsKey("answer")) {
        toret = ((Map<String, Object>) response.get("answer_box")).get("answer").toString();
      } else if (response.containsKey("answer_box")
        && ((Map<String, Object>) response.get("answer_box")).containsKey("snippet")) {
        toret = ((Map<String, Object>) response.get("answer_box")).get("snippet").toString();
      } else if (response.containsKey("answer_box")
        && ((Map<String, Object>) response.get("answer_box")).containsKey("snippet_highlighted_words")) {
        toret = ((List<String>) ((Map<String, Object>) response.get("answer_box"))
          .get("snippet_highlighted_words")).get(0);
      } else if (response.containsKey("sports_results")
        && ((Map<String, Object>) response.get("sports_results")).containsKey("game_spotlight")) {
        toret = ((Map<String, Object>) response.get("sports_results")).get("game_spotlight").toString();
      } else if (response.containsKey("shopping_results")
        && ((List<Map<String, Object>>) response.get("shopping_results")).get(0).containsKey("title")) {
        List<Map<String, Object>> shoppingResults = (List<Map<String, Object>>) response
          .get("shopping_results");
        List<Map<String, Object>> subList = shoppingResults.subList(0, 3);
        toret = subList.toString();
      } else if (response.containsKey("knowledge_graph")
        && ((Map<String, Object>) response.get("knowledge_graph")).containsKey("description")) {
        toret = ((Map<String, Object>) response.get("knowledge_graph")).get("description").toString();
      } else if ((((List<Map<String, Object>>) response.get("organic_results")).get(0)).containsKey("snippet")) {
        toret = (((List<Map<String, Object>>) response.get("organic_results")).get(0)).get("snippet")
          .toString();
      } else if ((((List<Map<String, Object>>) response.get("organic_results")).get(0)).containsKey("link")) {
        toret = (((List<Map<String, Object>>) response.get("organic_results")).get(0)).get("link").toString();
      } else if (response.containsKey("images_results")
        && ((Map<String, Object>) ((List<Map<String, Object>>) response.get("images_results")).get(0))
        .containsKey("thumbnail")) {
        List<String> thumbnails = new ArrayList<>();
        List<Map<String, Object>> imageResults = (List<Map<String, Object>>) response.get("images_results");
        for (Map<String, Object> item : imageResults.subList(0, 10)) {
          thumbnails.add(item.get("thumbnail").toString());
        }
        toret = thumbnails.toString();
      } else {
        toret = "No good search result found";
      }
      log.warn("SerpapiTool result:{}", toret);
      this.lastSearchResults = toret;
      return new ToolExecuteResult(toret);
    } catch (Exception e) {
      log.error("Error executing Google search", e);
      return new ToolExecuteResult("Error executing Google search: " + e.getMessage());
    }
  }

  @Override
  public String getCurrentToolStateString() {
    return String.format("""
        Google Search Status:
        - Search Location: %s
        - Recent Search: %s
        - Search Results: %s
        """, new java.io.File("").getAbsolutePath(),
      lastQuery.isEmpty() ? "No search performed yet"
        : String.format("Searched for: '%s' (max results: %d)", lastQuery, lastNumResults),
      lastSearchResults.isEmpty() ? "No results found" : lastSearchResults);
  }

  @Override
  public void cleanup(String planId) {
    // do nothing
  }

  @Override
  public String getServiceGroup() {
    return "default-service-group";
  }

  /**
   * Internal input class for defining input parameters of Google search tool
   */
  public static class GoogleSearchInput {

    private String query;

    @com.fasterxml.jackson.annotation.JsonProperty("num_results")
    private Integer numResults;

    public GoogleSearchInput() {
    }

    public GoogleSearchInput(String query, Integer numResults) {
      this.query = query;
      this.numResults = numResults;
    }

    public String getQuery() {
      return query;
    }

    public void setQuery(String query) {
      this.query = query;
    }

    public Integer getNumResults() {
      return numResults;
    }

    public void setNumResults(Integer numResults) {
      this.numResults = numResults;
    }

  }

  @Override
  public boolean isSelectable() {
    return true;
  }

}
